import numpy as np

def line_cylinder_intersection_z_axis(o, d, p0, r, h=None):
    """
    计算直线与轴线在 z 轴上的圆柱的交点。

    参数：
    - o: 直线上的一个点（三维 numpy 数组，形状为 (3,)）。
    - d: 直线的方向向量（三维 numpy 数组，形状为 (3,)）。
    - p0: 圆柱底面的中心（三维 numpy 数组，形状为 (3,)）。
    - r: 圆柱的半径。
    - h: 圆柱的高度（可选）。如果为 None，则圆柱无限长。

    返回：
    - 交点列表（每个交点为三维 numpy 数组，形状为 (3,)）。
    """
    # 提取 x 和 y 分量
    o_x, o_y, o_z = o[0], o[1], o[2]
    d_x, d_y, d_z = d[0], d[1], d[2]
    z0 = p0[2]  # 底面中心的 z 坐标

    # 二次方程的系数
    A = d_x**2 + d_y**2
    B = 2 * (o_x * d_x + o_y * d_y)
    C = o_x**2 + o_y**2 - r**2

    # 解二次方程
    discriminant = B**2 - 4 * A * C
    if discriminant < 0:
        return []  # 无实数解

    sqrt_discriminant = np.sqrt(discriminant)
    t1 = (-B + sqrt_discriminant) / (2 * A)
    t2 = (-B - sqrt_discriminant) / (2 * A)

    # 计算交点
    for t in [t1, t2]:
        point = o + t * d
        # 检查 z 坐标是否在 [z0, z0 + h] 范围内
        if z0 <= point[2] <= z0 + h:
            return True
    return False

# 示例用法
if __name__ == "__main__":
    # 定义直线：o + t * d
    o = np.array([-2, 0, 0])  # 直线经过原点
    d = np.array([1, 0, 1])  # 直线方向沿 x 轴

    # 定义圆柱：轴线在 z 轴上，底面中心在 (0, 0, 1)，半径为 1，高度为 2
    p0 = np.array([0, 0, 1])
    r = 1
    h = 2

    # 计算交点
    intersections = line_cylinder_intersection_z_axis(o, d, p0, r, h)
    print("是否相交：", intersections)
